source activate base # autodl
# source activate ke2torch23cu121 # 3090
export HUGGINGFACE_CACHE=/root/autodl-fs/huggingface/
export PYTHONUNBUFFERED=1

model=llama2-7b-chat
method=ICE

DATE=$(date +"%Y-%m-%d")
# autodl
LOG_ROOT="/root/autodl-fs/${method}-log/${DATE}"
LOG_DIRS=("${LOG_ROOT}" "${LOG_ROOT}/${method}-js" "${LOG_ROOT}/${method}-kl" "${LOG_ROOT}/${method}-wasserstein")

for log_dir in "${LOG_DIRS[@]}"; do
    if [ ! -d "${log_dir}" ]; then
        mkdir -p "${log_dir}"
    fi
done

# # KnowEdit datasets
# declare -A datasets=(
#     ["zsre"]="../dataset-ke/KnowEdit/benchmark/ZsRE/ZsRE-test-all.json"
#     ["recent"]="../dataset-ke/KnowEdit/benchmark/wiki_recent/recent_test.json"
#     ["counterfact"]="../dataset-ke/KnowEdit/benchmark/wiki_counterfact/test_cf.json"
#     ["wikibio"]="../dataset-ke/KnowEdit/benchmark/WikiBio/wikibio-test-all.json"
# )

gpu_id=0
# ICE KnowEdit datasets
declare -A datasets=(
    # ["zsre"]="../dataset-ke/${method}/zsre.json"
    # ["recent"]="../dataset-ke/${method}/wikidata_recent.json"
    # ["counterfact"]="../dataset-ke/${method}/wikidata_counterfact.json"
    # ["wikibio"]="../dataset-ke/${method}/wikibio.json"
)
# --test_generation 是否测试生成,计算fluency指标
# --save_gen_sentence 是否保存生成的句子,观察编辑后输出重复等问题
# 默认为false,指定位true
method_type=js
loss_div_alpha=1.0 # 主要是loss不下降,比例大小不影响.
early_stop_steps=5
layers=21
# module=origin
# for layers in 21 22,23,24,25,26 27,28,29,30,31 ; do

for module in mlp attn all; do
# for method_type in kl js; do
    cnt="fast-ly:${layers}-tg_ctx-${method_type}-${wst_type}:${loss_div_alpha}-stop:${early_stop_steps}-${module}"
    for datatype in "${!datasets[@]}"; do
        data_dir="${datasets[$datatype]}"
        echo "${DATE}/${method}-${method_type}/${model}-${datatype}-${gpu_id}-${cnt}"
        CUDA_VISIBLE_DEVICES=${gpu_id} python examples/run_knowedit_llama2.py \
            --editing_method=${method} \
            --layers=${layers} \
            --module ${module} \
            --loss_div_alpha=${loss_div_alpha} \
            --early_stop_steps=${early_stop_steps} \
            --p_dist=${method_type} \
            --datatype=${datatype} \
            --hparams_dir=./hparams/${method}/${model}.yaml \
            --objective_optimization=target_new_with_context \
            --data_dir=${data_dir} \
            --metrics_save_dir=/root/autodl-fs/${method}-results/${model}/${method}-${method_type}-${DATE}-${cnt} \
            --gen_sentence_save_dir=/root/autodl-fs/${method}-outputs/${model}/${method}-${method_type}-${DATE}-${cnt} \
            --pre_file=/root/autodl-fs/pre_edit/${model}_${datatype}_pre_edit.json \
            > /root/autodl-fs/${method}-log/${DATE}/${method}-${method_type}/${model}-${datatype}-${gpu_id}-${cnt}.log 2>&1 &
        wait
    done
done

wait && /usr/bin/shutdown
